
import numpy as np
from typing import Optional, List, Callable

class PSO:
    '''
    默认求解最大化 fitness 问题
    '''
    def __init__(self, D, fitness:Callable[[np.ndarray], float], N = 100, G = 50, v_bound = [-1, 1], x_low: Optional[List[int]] = None, x_high: Optional[List[int]] = None) -> None:
        self.D = D
        self.N = N
        self.G = G
        self.fitness = fitness
        self.fitness_list = []
        # 二维的约束矩阵 2 * D
        self.default_min = -10000
        self.default_max = 10000
        self.x_bound = np.zeros((2, D))
        self.v_bound = v_bound
        if x_low is not None:
            self.x_bound[0] = x_low
        # 没有约束情况
        else:
            self.x_bound[0] = np.full(D, self.default_min)
        if x_high is not None:
            self.x_bound[1] = x_high
        else:
            self.x_bound[1] = np.full(D, self.default_max)
        self.x = np.random.uniform(self.x_bound[0], self.x_bound[1], (N, D))
        self.v = np.random.uniform(v_bound[0], v_bound[1], (N, D))
        self.p_best = self.x.copy()
        self.g_best = self.p_best[np.argmax(self.fitness(self.p_best))]

    def update(self):
        w = 0.8  # 惯性权重
        c1 = 2.   # 个人学习因子
        c2 = 2.   # 社会学习因子

        # 对于种群中每一个个体:
        for i in range(self.N):
            r1 = np.random.rand(self.D) # 和 random.uniform(0, 1) 哪个好？
            r2 = np.random.rand(self.D)
            # 粒子群算法速度更新公式和位置更新公式
            self.v[i] = w * self.v[i] + c1 * r1 * (self.p_best[i] - self.x[i]) + c2 * r2 * (self.g_best - self.x[i])
            self.x[i] += self.v[i]
            # 位置和速度限制
            self.x[i] = np.clip(self.x[i], self.x_bound[0], self.x_bound[1])
            self.v[i] = np.clip(self.v[i], self.v_bound[0], self.v_bound[1])
            # 更新个体最优和全局最优
            if self.fitness(self.x[i]) > self.fitness(self.p_best[i]):
                self.p_best[i] = self.x[i]
            if self.fitness(self.p_best[i]) > self.fitness(self.g_best):
                self.g_best = self.p_best[i]
        self.fitness_list.append(self.fitness(self.g_best))

    def pso(self):
        for i in range(self.G):
            self.update()
        return self.g_best, self.fitness(self.g_best)

    def draw_fitness_history(self):
        import matplotlib.pyplot as plt
        idx = [i for i in range(len(self.fitness_list))]
        plt.plot(idx, self.fitness_list)
        plt.xlabel('Iteration')
        plt.ylabel('Fitness')
        plt.title('Fitness History')
        plt.show()

if __name__ == '__main__':

    pso = PSO(5, lambda x: -(x[0]-1)**2 - (x[1]-2.5)**2 - (x[2]-3)**2 - (x[3]-4)**2 - (x[4]-5)**2, N=100, G=200, v_bound=[-1, 1])
    pso.pso()
    pso.draw_fitness_history()